"""
使用TVDI的另一种降尺度算法
土壤和植被分量温度（SVCT）
TVDI 与土壤湿度之间的负线性关系
"""
import os

import numpy as np
from osgeo import gdal, osr
from scipy.optimize import leastsq, curve_fit
from tqdm import tqdm


# 写入tif文件
def write_tif(file_path, data, geotrans, projection, nodata, gdal_type):
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    dataset.SetGeoTransform(geotrans)
    # 定义投影
    # prj = osr.SpatialReference()
    # prj.ImportFromEPSG(4326)
    # dataset.SetProjection(prj.ExportToWkt())
    dataset.SetProjection(projection)
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    dataset.FlushCache()
    print("写入成功：{}".format(file_path))


# 获取低分辨率像元对应所有高分辨率像元值
def get_high_res_pixels(high_res_geo_transform, high_res_data, low_res_geo_transform, x_low_res_pixel, y_low_res_pixel, low_res_to_high_res_ratio):
    # 计算低分辨率影像像元在高分辨率影像中的范围
    x_low_res_geo, y_low_res_geo = pixel2coord(low_res_geo_transform, x_low_res_pixel, y_low_res_pixel)
    x_low_res_geo = x_low_res_geo + low_res_geo_transform[1] / 2  # 左右偏移
    y_low_res_geo = y_low_res_geo + low_res_geo_transform[5] / 2  # 上下偏移

    # 计算行列号范围
    high_res_x, high_res_y = coord2pixel(high_res_geo_transform, x_low_res_geo, y_low_res_geo)

    high_res_x_min = int(high_res_x - low_res_to_high_res_ratio / 2)
    high_res_x_max = int(high_res_x + low_res_to_high_res_ratio / 2)
    high_res_y_min = int(high_res_y - low_res_to_high_res_ratio / 2)
    high_res_y_max = int(high_res_y + low_res_to_high_res_ratio / 2)

    # 排除边缘像元的影响
    if high_res_x_min < 0 and high_res_x_max > 0:
        high_res_x_min = 0
    if high_res_x_min < high_res_data.shape[1] < high_res_x_max:
        high_res_x_max = high_res_data.shape[1]
    if high_res_y_min < 0 and high_res_y_max > 0:
        high_res_y_min = 0
    if high_res_y_min < high_res_data.shape[0] < high_res_y_max:
        high_res_y_max = high_res_data.shape[0]

    # 获取范围内的像元值
    if 0 <= high_res_x_min <= high_res_data.shape[1] and 0 <= high_res_y_min <= high_res_data.shape[0]:
        high_res_pixels = high_res_data[high_res_y_min:high_res_y_max, high_res_x_min:high_res_x_max]
        return high_res_pixels.astype(np.float32)
    else:
        return []


# 根据像元行列号计算经纬度
def pixel2coord(geo_transform, x, y):
    x_origin, pixel_width, x_rotation, y_origin, y_rotation, pixel_height = geo_transform
    x_geo = x_origin + pixel_width * x
    y_geo = y_origin + pixel_height * y
    return x_geo, y_geo


# 根据经纬度计算像元行列号
def coord2pixel(geo_transform, x_geo, y_geo):
    x = int((x_geo - geo_transform[0]) / geo_transform[1])
    y = int((y_geo - geo_transform[3]) / geo_transform[5])
    return x, y


# 计算SVCT(土壤和植被成分温度)
def svct(lst, fc):
    """
    两个假设来求解SVCT
    假设1：相邻像素的 SVCT 相近。
    假设2：同一像素点的植被成分温度低于土壤成分温度。
    选取符合假设2的像素点及相邻的8个像素点，根据方程建立方程，通过最小二乘法近似求解方程得到SVCT
    :param ndvi_path:
    :param lst_path:
    :return:
    """
    Ts_array = np.ones_like(lst) * -9999
    Tv_array = np.ones_like(lst) * -9999
    epsilons = 0.97
    epsilonv = 0.985
    delta = 5.670373e-8  # 斯蒂芬-玻尔兹曼常数

    # print(ndvi.min())

    # epsilon = 0  # 发射率
    # T = 0  # 表面温度
    epsilon = (1 - fc) * epsilons + fc * epsilonv

    # 拟合等式
    def func(x, p):
        epsilons = 0.97
        epsilonv = 0.985
        delta = 5.670373e-8  # 斯蒂芬-玻尔兹曼常数
        fc = x
        Ts, Tv = p
        return (1 - fc) * delta * epsilons * Ts ** 4 + fc * delta * epsilonv * Tv ** 4

    # 拟合误差函数
    def residuals(p, y, x):
        return y - func(x, p)

    for i in range(lst.shape[0]):
        for j in range(lst.shape[1]):
            fc_neibor = []
            T_neribor = []
            epsilon_neribor = []
            # 不在边缘
            if i != 0 and j != 0 and i != (lst.shape[0] - 1) and j != (lst.shape[1] - 1):
                fc_neibor = np.array([fc[i, j], fc[i - 1, j - 1], fc[i - 1, j], fc[i - 1, j + 1], fc[i, j - 1], fc[i, j + 1], fc[i + 1, j - 1], fc[i + 1, j], fc[i + 1, j + 1]])
                T_neribor = np.array([lst[i, j], lst[i - 1, j - 1], lst[i - 1, j], lst[i - 1, j + 1], lst[i, j - 1], lst[i, j + 1], lst[i + 1, j - 1], lst[i + 1, j], lst[i + 1, j + 1]])
                epsilon_neribor = np.array(
                    [epsilon[i, j], epsilon[i - 1, j - 1], epsilon[i - 1, j], epsilon[i - 1, j + 1], epsilon[i, j - 1], epsilon[i, j + 1], epsilon[i + 1, j - 1], epsilon[i + 1, j],
                     epsilon[i + 1, j + 1]])
            # 在上边缘且不在顶点
            elif i == 0 and j != 0 and j != (lst.shape[1] - 1):
                fc_neibor = np.array([fc[i, j], fc[i, j - 1], fc[i, j + 1], fc[i + 1, j - 1], fc[i + 1, j], fc[i + 1, j + 1]])
                T_neribor = np.array([lst[i, j], lst[i, j - 1], lst[i, j + 1], lst[i + 1, j - 1], lst[i + 1, j], lst[i + 1, j + 1]])
                epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j - 1], epsilon[i, j + 1], epsilon[i + 1, j - 1], epsilon[i + 1, j], epsilon[i + 1, j + 1]])
            # 在左边缘且不在顶点
            elif i != 0 and j == 0 and i != (lst.shape[0] - 1):
                fc_neibor = np.array([fc[i, j], fc[i - 1, j], fc[i + 1, j], fc[i, j + 1], fc[i - 1, j + 1], fc[i + 1, j + 1]])
                T_neribor = np.array([lst[i, j], lst[i - 1, j], lst[i + 1, j], lst[i, j + 1], lst[i - 1, j + 1], lst[i + 1, j + 1]])
                epsilon_neribor = np.array([epsilon[i, j], epsilon[i - 1, j], epsilon[i + 1, j], epsilon[i, j + 1], epsilon[i - 1, j + 1], epsilon[i + 1, j + 1]])
            # 在下边缘且不在顶点
            elif i == (lst.shape[0] - 1) and j != 0 and j != (lst.shape[1] - 1):
                fc_neibor = np.array([fc[i, j], fc[i, j - 1], fc[i, j + 1], fc[i - 1, j - 1], fc[i - 1, j], fc[i - 1, j + 1]])
                T_neribor = np.array([lst[i, j], lst[i, j - 1], lst[i, j + 1], lst[i - 1, j - 1], lst[i - 1, j], lst[i - 1, j + 1]])
                epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j - 1], epsilon[i, j + 1], epsilon[i - 1, j - 1], epsilon[i - 1, j], epsilon[i - 1, j + 1]])
            # 在右边缘且不在顶点
            elif i != 0 and j == (lst.shape[1] - 1) and i != (lst.shape[0] - 1):
                fc_neibor = np.array([fc[i, j], fc[i - 1, j], fc[i + 1, j], fc[i, j - 1], fc[i - 1, j - 1], fc[i + 1, j - 1]])
                T_neribor = np.array([lst[i, j], lst[i - 1, j], lst[i + 1, j], lst[i, j - 1], lst[i - 1, j - 1], lst[i + 1, j - 1]])
                epsilon_neribor = np.array([epsilon[i, j], epsilon[i - 1, j], epsilon[i + 1, j], epsilon[i, j - 1], epsilon[i - 1, j - 1], epsilon[i + 1, j - 1]])
            else:
                if i == 0 and j == 0:
                    fc_neibor = np.array([fc[i, j], fc[i, j + 1], fc[i + 1, j], fc[i + 1, j + 1]])
                    T_neribor = np.array([lst[i, j], lst[i, j + 1], lst[i + 1, j], lst[i + 1, j + 1]])
                    epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j + 1], epsilon[i + 1, j], epsilon[i + 1, j + 1]])
                elif i == 0 and j == (lst.shape[1] - 1):
                    fc_neibor = np.array([fc[i, j], fc[i, j - 1], fc[i + 1, j], fc[i + 1, j - 1]])
                    T_neribor = np.array([lst[i, j], lst[i, j - 1], lst[i + 1, j], lst[i + 1, j - 1]])
                    epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j - 1], epsilon[i + 1, j], epsilon[i + 1, j - 1]])
                elif i == (lst.shape[0] - 1) and j == 0:
                    fc_neibor = np.array([fc[i, j], fc[i, j + 1], fc[i - 1, j], fc[i - 1, j + 1]])
                    T_neribor = np.array([lst[i, j], lst[i, j + 1], lst[i - 1, j], lst[i - 1, j + 1]])
                    epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j + 1], epsilon[i - 1, j], epsilon[i - 1, j + 1]])
                elif i == (lst.shape[0] - 1) and j == (lst.shape[1] - 1):
                    fc_neibor = np.array([fc[i, j], fc[i, j - 1], fc[i - 1, j], fc[i - 1, j - 1]])
                    T_neribor = np.array([lst[i, j], lst[i, j - 1], lst[i - 1, j], lst[i - 1, j - 1]])
                    epsilon_neribor = np.array([epsilon[i, j], epsilon[i, j - 1], epsilon[i - 1, j], epsilon[i - 1, j - 1]])
            x = fc_neibor[np.isnan(fc_neibor) == False]
            epsilon_new = epsilon_neribor[np.isnan(fc_neibor) == False]
            y = T_neribor[np.isnan(fc_neibor) == False] ** 4 * delta * epsilon_new
            if len(x) < 2:
                Ts_array[i, j] = np.nan
                Tv_array[i, j] = np.nan
                continue
            p0 = np.array([100, 80])
            plsq = leastsq(residuals, p0, args=(y, x))
            Ts, Tv = plsq[0]
            Ts_array[i, j] = Ts
            Tv_array[i, j] = Tv
    return Ts_array, Tv_array


# 降尺度公式
def downscale_formula(x, a, c, m, n):
    Ts, Tv, fc = x.T
    # a_x = a / (Ts.max() - Ts.min())
    # c_x = c / (Tv.max() - Tv.min())
    # m = a_x * Ts.min() - c_x * Tv.min() - b + d
    # n = b - a_x * Ts.min()
    # print((1 - fc) * Ts)
    # print(fc * Tv)
    # print(fc)
    sm = a * (1 - fc) * Ts + c * fc * Tv + m * fc + n
    return sm


def fitting_formula(x, a, c, m, n):
    Ts, Tv, fc = x.T
    # print(np.nanmean((1 - fc) * Ts))
    # print(np.nanmean(fc * Tv))
    # print(np.nanmean(fc))
    sml = a * np.nanmean((1 - fc) * Ts) + c * np.nanmean(fc * Tv) + m * np.nanmean(fc) + n
    return sml


def main(ndvi_dir, lst_dir, tdvi_dir, sm_dir, output_dir):
    ndvi_list = [os.path.join(ndvi_dir, file) for file in os.listdir(ndvi_dir) if file.endswith(".tif")]
    lst_list = [os.path.join(lst_dir, file) for file in os.listdir(lst_dir) if file.endswith(".tif")]
    sm_list = [os.path.join(sm_dir, file) for file in os.listdir(sm_dir) if file.endswith(".tif")]
    for i in tqdm(range(len(ndvi_list)), desc="土壤水分降尺度"):
        ndvi_path = ndvi_list[i]
        lst_path = lst_list[i]
        sm_path = sm_list[i]

        # print(Ts, Tv)
        ndvi_dataset = gdal.Open(ndvi_path)
        lst_dataset = gdal.Open(lst_path)
        sm_dataset = gdal.Open(sm_path)

        ndvi = ndvi_dataset.ReadAsArray()
        lst = lst_dataset.ReadAsArray()
        ndvi[ndvi == -9999] = np.nan
        lst[lst == -9999] = np.nan
        ndvi = ndvi * 0.0001
        lst = lst * 0.02
        # print(np.nanmin(ndvi))
        # print(np.nanmin(lst))
        sm = sm_dataset.ReadAsArray()

        high_geotrans = ndvi_dataset.GetGeoTransform()
        low_geotrans = sm_dataset.GetGeoTransform()

        high_projection = ndvi_dataset.GetProjection()

        # 像元比例
        res_ratio = int(low_geotrans[1] / high_geotrans[1])
        # print(res_ratio)  # 30

        # 植被覆盖度由ndvi计算,这里需要确定两个置信区间
        # 以前0.5%的值为土地覆盖ndvi，以99.5%的值为植被全覆盖ndvi，这个区间都是经验性的。
        # 计算公式为fc = (NDVI - NDVIsoil) / (NDVIveg - NDVIsoil)
        ndvi_new = ndvi[ndvi >= 0]
        NDVIsoil = np.percentile(ndvi_new, 0.5)
        NDVIveg = np.percentile(ndvi_new, 99.5)
        # ndvi[ndvi == -0.9999] = np.nan
        # ndvi[ndvi < NDVIsoil] = np.nan
        # print(np.nanmin(ndvi))
        fc = (ndvi - NDVIsoil) / (NDVIveg - NDVIsoil)
        fc[fc < 0] = 0
        fc[fc > 1] = 1
        # print(np.nanmin(lst))
        # print(np.nanmax(lst))
        # print(np.nanmin(fc))
        # print(np.nanmax(fc))
        # print("=" * 10)
        Ts, Tv = svct(lst, fc)
        # Ts[Ts < 0] = np.nan
        # Tv[Tv < 0] = np.nan
        # print(np.nanmin(Ts))
        # print(np.nanmin(Tv))

        # print("="*10)

        # 拟合部分
        Ts_in_sm = []
        Tv_in_sm = []
        fc_in_sm = []
        for y in range(sm.shape[0]):
            for x in range(sm.shape[1]):
                Ts_in_sm.append(get_high_res_pixels(high_geotrans, Ts, low_geotrans, x, y, res_ratio))
                Tv_in_sm.append(get_high_res_pixels(high_geotrans, Tv, low_geotrans, x, y, res_ratio))
                fc_in_sm.append(get_high_res_pixels(high_geotrans, fc, low_geotrans, x, y, res_ratio))
        Ts_in_sm = np.asarray(Ts_in_sm, dtype=np.float32)
        Tv_in_sm = np.asarray(Tv_in_sm, dtype=np.float32)
        fc_in_sm = np.asarray(fc_in_sm, dtype=np.float32)
        sm_flatten = sm.flatten()
        # print(sm_flatten)
        Ts_filter = Ts_in_sm[sm_flatten != -9999]
        Tv_filter = Tv_in_sm[sm_flatten != -9999]
        fc_filter = fc_in_sm[sm_flatten != -9999]
        sm_filter = sm_flatten[sm_flatten != -9999]
        if len(sm_filter) < 4:
            print("{}SM可用数据少于4个无法拟合".format(os.path.splitext(ndvi_path.split("\\")[-1])[0].split("_")[0]))
            continue

        # 将异常、缺失值全部转为nan
        # Ts_filter_new = [np.where(arr == -9999, np.nan, arr) for arr in Ts_filter]
        Ts_filter_new = [np.where(arr == np.inf, np.nan, arr) for arr in Ts_filter]
        # Tv_filter_new = [np.where(arr == -9999, np.nan, arr) for arr in Tv_filter]
        Tv_filter_new = [np.where(arr == np.inf, np.nan, arr) for arr in Tv_filter]
        # fc_filter_new = [np.where(arr == -9999, np.nan, arr) for arr in fc_filter]
        fc_filter_new = [np.where(arr == np.inf, np.nan, arr) for arr in fc_filter]
        # print("拟合输入")
        # print(np.nanmin(Ts_filter_new))
        # print(np.nanmin(Tv_filter_new))
        # print(np.nanmin(fc_filter_new))
        # print(np.nanmin(sm_filter))
        # print(np.nanmax(Ts_filter_new))
        # print(np.nanmax(Tv_filter_new))
        # print(np.nanmax(fc_filter_new))
        # print(np.nanmax(sm_filter))
        # print("="*10)

        X = np.vstack((np.concatenate([arr.flatten() for arr in Ts_filter_new]),
                       np.concatenate([arr.flatten() for arr in Tv_filter_new]),
                       np.concatenate([arr.flatten() for arr in fc_filter_new]))).T
        # print(sm_filter)
        popt, pcov = curve_fit(fitting_formula, X, sm_filter, maxfev=1000000, check_finite=False, bounds=(-2, 2))
        print("拟合系数：", popt)
        a, c, m, n = popt[0], popt[1], popt[2], popt[3]
        # print(a, c, m, n)

        # 降尺度
        Ts_high = np.where(Ts == -9999, np.nan, Ts)
        Tv_high = np.where(Tv == -9999, np.nan, Tv)
        fc_high = np.where(fc == -9999, np.nan, fc)
        Ts_high[Ts_high < 0] = np.nan
        Tv_high[Tv_high < 0] = np.nan
        # print("降尺度输入")
        # print(np.nanmin(Ts_high))
        # print(np.nanmin(Tv_high))
        # print(np.nanmin(fc_high))
        # print(np.nanmax(Ts_high))
        # print(np.nanmax(Tv_high))
        # print(np.nanmax(fc_high))
        # print("=" * 10)

        X_high = np.vstack((Ts_high.flatten(), Tv_high.flatten(), fc_high.flatten())).T
        sm_high = downscale_formula(X_high, a, c, m, n)
        sm_high = sm_high.reshape(ndvi.shape)
        # print(np.nanmin(sm_high))
        # print(np.nanmax(sm_high))
        sm_high[np.isnan(sm_high)] = -9999
        output_path = output_dir + "/" + os.path.splitext(ndvi_path.split("\\")[-1])[0].split("_")[0] + ".tif"
        write_tif(output_path, sm_high, high_geotrans, high_projection, -9999, gdal.GDT_Float32)


if __name__ == '__main__':
    ndvi_dir = r"G:\test\Mask\ndvi_filter_mask"
    lst_dir = r"G:\test\Mask\day_lst_mask"
    tvdi_dir = r"G:\test\TDVI"
    sm_dir = r"G:\test\sm_cut"
    output_dir = r"G:\test\downscal_new"
    main(ndvi_dir, lst_dir, tvdi_dir, sm_dir, output_dir)
